// Copyright (C) Kumo inc. and its affiliates.
// Author: Jeff.li lijippy@163.com
// All rights reserved.
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published
// by the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.
//


#include <utility>

#include <melon/hash/checksum.h>
#include <pollux/common/file/file_input_stream.h>
#include <pollux/exec/operator_trace_reader.h>
#include <pollux/exec/trace_util.h>

namespace kumo::pollux::exec::trace {
    OperatorTraceInputReader::OperatorTraceInputReader(
        std::string traceDir,
        RowTypePtr dataType,
        memory::MemoryPool *pool)
        : traceDir_(std::move(traceDir)),
          fs_(filesystems::getFileSystem(traceDir_, nullptr)),
          dataType_(std::move(dataType)),
          pool_(pool),
          serde_(getNamedVectorSerde(VectorSerde::Kind::kPresto)),
          inputStream_(getInputStream()) {
        POLLUX_CHECK_NOT_NULL(dataType_);
    }

    bool OperatorTraceInputReader::read(RowVectorPtr &batch) const {
        if (inputStream_ == nullptr) {
            return false;
        }

        if (inputStream_->atEnd()) {
            batch = nullptr;
            return false;
        }

        VectorStreamGroup::read(
            inputStream_.get(), pool_, dataType_, serde_, &batch, &readOptions_);
        return true;
    }

    std::unique_ptr<common::FileInputStream>
    OperatorTraceInputReader::getInputStream() const {
        auto traceFile = fs_->openFileForRead(getOpTraceInputFilePath(traceDir_));
        if (traceFile->size() == 0) {
            KLOG(WARNING) << "Operator trace input file is empty: "
                 << getOpTraceInputFilePath(traceDir_);
            return nullptr;
        }
        // TODO: Make the buffer size configurable.
        return std::make_unique<common::FileInputStream>(
            std::move(traceFile), 1 << 20, pool_);
    }

    OperatorTraceSummaryReader::OperatorTraceSummaryReader(
        std::string traceDir,
        memory::MemoryPool * /* pool */)
        : traceDir_(std::move(traceDir)),
          fs_(filesystems::getFileSystem(traceDir_, nullptr)),
          summaryFile_(fs_->openFileForRead(getOpTraceSummaryFilePath(traceDir_))) {
    }

    OperatorTraceSummary OperatorTraceSummaryReader::read() const {
        POLLUX_CHECK_NOT_NULL(summaryFile_);
        const auto summaryStr = summaryFile_->pread(0, summaryFile_->size());
        POLLUX_CHECK(!summaryStr.empty());

        melon::Dynamic summaryObj = melon::parseJson(summaryStr);
        OperatorTraceSummary summary;
        summary.opType = summaryObj[OperatorTraceTraits::kOpTypeKey].asString();
        if (summary.opType == "TableScan") {
            summary.numSplits = summaryObj[OperatorTraceTraits::kNumSplitsKey].asInt();
        }
        summary.peakMemory = summaryObj[OperatorTraceTraits::kPeakMemoryKey].asInt();
        summary.inputRows = summaryObj[OperatorTraceTraits::kInputRowsKey].asInt();
        summary.inputBytes = summaryObj[OperatorTraceTraits::kInputBytesKey].asInt();
        summary.rawInputRows =
                summaryObj[OperatorTraceTraits::kRawInputRowsKey].asInt();
        summary.rawInputBytes =
                summaryObj[OperatorTraceTraits::kRawInputBytesKey].asInt();
        return summary;
    }

    OperatorTraceSplitReader::OperatorTraceSplitReader(
        std::vector<std::string> traceDirs,
        memory::MemoryPool *pool)
        : traceDirs_(std::move(traceDirs)),
          fs_(filesystems::getFileSystem(traceDirs_[0], nullptr)),
          pool_(pool) {
        POLLUX_CHECK_NOT_NULL(fs_);
    }

    std::vector<std::string> OperatorTraceSplitReader::read() const {
        std::vector<std::string> splits;
        for (const auto &traceDir: traceDirs_) {
            auto stream = getSplitInputStream(traceDir);
            if (stream == nullptr) {
                continue;
            }
            auto curSplits = deserialize(stream.get());
            splits.insert(
                splits.end(),
                std::make_move_iterator(curSplits.begin()),
                std::make_move_iterator(curSplits.end()));
        }
        return splits;
    }

    std::unique_ptr<common::FileInputStream>
    OperatorTraceSplitReader::getSplitInputStream(
        const std::string &traceDir) const {
        auto splitInfoFile = fs_->openFileForRead(getOpTraceSplitFilePath(traceDir));
        if (splitInfoFile->size() == 0) {
            KLOG(WARNING) << "Split info is empty in " << traceDir;
            return nullptr;
        }
        // TODO: Make the buffer size configurable.
        return std::make_unique<common::FileInputStream>(
            std::move(splitInfoFile), 1 << 20, pool_);
    }

    // static
    std::vector<std::string> OperatorTraceSplitReader::deserialize(
        common::FileInputStream *stream) {
        std::vector<std::string> splits;
        try {
            while (!stream->atEnd()) {
                const auto length = stream->read<uint32_t>();
                std::string splitStr(length, '\0');
                stream->readBytes(reinterpret_cast<uint8_t *>(splitStr.data()), length);
                const auto crc32 = stream->read<uint32_t>();
                const auto actualCrc32 = melon::crc32(
                    reinterpret_cast<const uint8_t *>(splitStr.data()), splitStr.size());
                if (crc32 != actualCrc32) {
                    KLOG(ERROR) << "Failed to verify the split checksum " << crc32
                   << " which does not equal to the actual computed checksum "
                   << actualCrc32;
                    break;
                }
                splits.push_back(std::move(splitStr));
            }
        } catch (const PolluxException &e) {
            KLOG(ERROR) << "Failed to deserialize split: " << e.message();
        }
        return splits;
    }
} // namespace kumo::pollux::exec::trace
